#!/usr/bin/perl

# Author: Lipyeow Lim
# Date  : 12/21/2010

use IO::Socket;
use threads;
use Data::Dumper;
use DBI;
use Log::Fast;
use Time::HiRes qw( usleep ualarm gettimeofday tv_interval nanosleep
                      clock_gettime clock_getres clock_nanosleep clock
                      stat );
use strict;

scalar @ARGV == 1 or 
   die "workerd.pl <port>\n";

my $workerhostport = $ARGV[0];

my $sqlite3ext= "/home/ubuntu/ps/data/libsqlitefunctions.so";

my $operators = {
   'DBEnrich' => \&dbenrich,
   'DBSource' => \&dbsource
};

#my ($myid) = split(/\./, `echo \$HOSTNAME`);
my $myid = `echo \$HOSTNAME`; chomp($myid);
$myid .= ":$workerhostport";

my $log = Log::Fast->new({
         level => 'DEBUG',
         prefix => $myid . ' %L: ',
         type => 'fh',
         fh => \*STDERR,
        });

my $sock = new IO::Socket::INET ( 
                  LocalHost => $workerhostport,
                  Proto => 'tcp',
                  Listen => 1,
                  ReuseAddr => 1);

die "Could not create socket : $!\n" unless $sock;

my $clientsock;
while ( $clientsock = $sock->accept())
{
   my $t = threads->new( \&workerthread, $clientsock );
   $t->detach();
}
close $sock;


#------------------------------------------------
#------------------------------------------------
sub workerthread {
   my ($clientsock) = @_;
   my $tid = threads->tid();
   $log->DEBUG("worker thread $tid executing ...\n");

   # read message into lines
   my @lines = ();
   while ( <$clientsock> )
   {
#      print "$tid: recv ", $_;
      chomp;
      push @lines, $_;
   }

   # parse message
   my $msg = parseWorkMessage(join("\n",@lines));

   #print Dumper($msg);
   my ($op) = $msg->{'cmd'}; 
   if(! exists $operators->{$op}) 
   {
      $log->ERR("$tid: Unrecognized Operator $op\n"); 
      return;
   }


   my ($insrvsocklist, $insocklist, $outsocklist) 
         = ([],[],[]);

   if ( setupInputSockets($msg->{'input'},
           $insrvsocklist,$insocklist) == 0 )
   {
      return;
   }

   if ( setupOutputSockets($msg->{'output'}, 
           $outsocklist) == 0 )
   {
      return;
   }

   $msg->{'insocklist'} = $insocklist;
   $msg->{'insrvsocklist'} = $insrvsocklist;
   $msg->{'outsocklist'} = $outsocklist;

   &{$operators->{$op}}($msg);
   $log->DEBUG( "$tid: $op\n");

   closeSockets( $outsocklist );
   closeSockets( $insocklist );
   closeSockets( $insrvsocklist );

   $log->DEBUG( "$tid: done\n");
}

#------------------------------------------------
# parse a message string
#------------------------------------------------
sub parseWorkMessage
{
   my ($msgstr) = @_;

   my $msg;
   
   eval $msgstr;

   # TODO: check command validity etc
   #print Dumper( $msg);

   return $msg;
}

#------------------------------------------------
# setup input sockets
#------------------------------------------------
sub setupInputSockets
{
   my ($hostlist, $insrvsock, $insock) = @_;

   foreach my $hostport (@$hostlist)
   {
      my $sock = new IO::Socket::INET ( LocalHost => $hostport,
                                  Proto => 'tcp',
                                  Listen => 1,
                                  ReuseAddr => 1
                          );
      if( ! $sock )
      {
          $log->Err("Could not create socket $hostport: $!\n");
          return 0;
      }

      $log->DEBUG( "Start listening on input stream
              $hostport\n");
      my $clientsock = $sock->accept();
      push @$insrvsock, $sock;
      push @$insock, $clientsock;
   }
   return 1;
}

#------------------------------------------------
# setup output sockets
#------------------------------------------------
sub setupOutputSockets 
{
   my ($hostlist, $outsock) = @_;

   my $sock;
   foreach my $hostport (@$hostlist)
   {
      my $tries = 1;
      my $backoff = 10;
      $sock = new IO::Socket::INET ( PeerAddr => $hostport,
                                  Proto => 'tcp');
      while(! $sock && $tries < 50 )
      {
          usleep($backoff);
          $sock = new IO::Socket::INET ( PeerAddr => $hostport,
                                  Proto => 'tcp');
          $tries++;
          $backoff *= 2;
      }
      if(! $sock )
      {
         $log->Err("Could not create socket $hostport: $!\n");
         return 0;
      }
      push @$outsock, $sock;
      $log->DEBUG( "Connected output stream to $hostport\n");
   }
   return 1;
}

#------------------------------------------------
#------------------------------------------------
sub closeSockets
{
   my ($socklist) = @_;

   foreach (@$socklist)
   {
      close $_;
   }
}

#------------------------------------------------
# loads a list of sqlite extension libraries
# returns 0 if error
#------------------------------------------------
sub sqlite_load_extensions
{
   my ($dbh,$extlist) = @_;

   scalar @$extlist > 0 or return 1;

   $dbh->sqlite_enable_load_extension(1);
   foreach my $extpath (@$extlist)
   {
      my $stmt = $dbh->prepare(
                  "select load_extension('$extpath')");
      if( ! defined $stmt )
      { 
         $log->ERR( "Cannot prepare: " . $dbh->errstr ); 
         return 0;
      }
   
      if ( ! $stmt->execute())
      { 
         $log->ERR($stmt->errstr);
         return 0;
      }
   }
   return 1;
}

#------------------------------------------------
#------------------------------------------------
sub dbsource
{
   my ($msg) = @_;

   my ($dbfile,$sql) = @{$msg->{'args'}};
   my $outsocklist = $msg->{'outsocklist'};

   my $dbh = DBI->connect("dbi:SQLite:dbname=$dbfile","","");
   if( ! defined $dbh )
   {
      $log->ERR("Cannot connect to $dbfile : " . $DBI::errstr);
      return 0;
   }
   #$log->DEBUG("Connected to $dbfile");

   if( exists $msg->{'sqlite_extensions'})
   {
      if( sqlite_load_extensions($dbh,
                      $msg->{'sqlite_extensions'}) == 0 )
      { 
         $log->ERR( "load extensions failed"); 
         return 0;
      }
      #$log->DEBUG("Loaded extensions ");
   }
   my $stmt = $dbh->prepare("$sql");
   if ( ! defined $stmt )
   {
      $log->ERR("Can't prepare : $DBI::errstr\n$sql\n");
      return 0;
   }

   if ( ! defined $stmt->execute() )
   {
      $log->ERR($stmt->errstr);
      return 0;
   }

   #$log->DEBUG("Executing $sql");

   my @row = $stmt->fetchrow_array;
   while( scalar @row > 0 )
   {
      my $rowstr = join("|", @row) . "\n";
      foreach my $outsock (@$outsocklist)
      {
         print $outsock $rowstr;
      }
      @row = $stmt->fetchrow_array;
   }

   $dbh->disconnect;
   return 1;
}

#------------------------------------------------
#------------------------------------------------
sub dbenrich
{
   my ($msg) = @_;

   my ($dbfile,$sql,$params) = @{$msg->{'args'}};

   my @argidx = split(/\,/,$params);

   # print "dbenrich: argidx = @argidx\n";
   my $outsocklist = $msg->{'outsocklist'};
   my $insocklist = $msg->{'insocklist'};
   my $insock = $insocklist->[0];

   my $dbh = DBI->connect("dbi:SQLite:dbname=$dbfile","","");
   if( ! defined $dbh )
   {
      $log->ERR("Cannot connect to $dbfile : " . $DBI::errstr);
      return 0;
   }
   #$log->DEBUG("Connected to $dbfile");

   if( exists $msg->{'sqlite_extensions'})
   {
      if( sqlite_load_extensions($dbh,
                      $msg->{'sqlite_extensions'}) == 0 )
      { 
         $log->ERR( "load extensions failed"); 
         return 0;
      }
      #$log->DEBUG("Loaded extensions ");
   }

   #$log->DEBUG( "prepare $sql\n");
   my $stmt = $dbh->prepare("$sql");
   if( ! defined $stmt )
   {
      $log->ERR("Can't prepare : $DBI::errstr\n$sql\n");
      return 0;
   }
   while( <$insock> )
   {
      chomp;
      my @intuple = split(/|/);

      #print "dbenrich: in $_\n";
      # bind all join parameters from the current tuple
      foreach my $i (@argidx)
      {
         #print "dbenrich: bind param $i+1 to $intuple[$i]\n";
         $stmt->bind_param($i+1, $intuple[$i]);
      }
      # execute the join of current tuple with table
      if( ! defined $stmt->execute())
      {  
         $log->ERR($stmt->errstr);
         return 0;
      }
      # read all rows from resultset
      my @row = $stmt->fetchrow_array;
      while( scalar @row > 0 )
      {
         my $rowstr = join("|", @row) . "\n";
         #print "dbenrich: join $rowstr";
         foreach my $outsock (@$outsocklist)
         {
            print $outsock "$_|$rowstr";
            #print "dbenrich: out $_|$rowstr";
         }
         @row = $stmt->fetchrow_array;
      }
   }
   
   $dbh->disconnect;
}
